import torch
import torchvision
import torchaudio

print("=== PyTorch Environment Check ===")
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"torchvision: {torchvision.__version__}")
print(f"torchaudio: {torchaudio.__version__}")

if torch.cuda.is_available():
    print(f"\nCUDA Information:")
    print(f"CUDA version: {torch.version.cuda}")
    print(f"Current device: {torch.cuda.current_device()}")
    print(f"Device name: {torch.cuda.get_device_name()}")
    print(f"Device memory: {torch.cuda.get_device_properties(0).total_memory / 1024**2:.0f}MB") 